package cn.linkai.star.graph.core.service.impl;

import cn.hutool.core.bean.BeanUtil;
import cn.hutool.core.util.StrUtil;
import cn.linkai.star.graph.comfyui.client.pojo.ComfyuiModel;
import cn.linkai.star.graph.comfyui.client.pojo.ComfyuiRequestDto;
import cn.linkai.star.graph.comfyui.client.pojo.ComfyuiTask;
import cn.linkai.star.graph.core.dto.request.Text2ImageReqDto;
import cn.linkai.star.graph.core.dto.respone.Text2ImageResDto;
import cn.linkai.star.graph.core.service.*;
import cn.linkai.star.graph.core.utils.UserUtils;
import com.alibaba.fastjson2.JSON;
import freemarker.template.TemplateException;
import lombok.RequiredArgsConstructor;
import org.springframework.stereotype.Service;

import java.io.IOException;

@RequiredArgsConstructor
@Service
public class Text2lmageServiceImpl implements Text2ImageService {

    public static final String COMFYUI_CLIENT_ID = "star-graph";

    //注入
    private final RedisService redisService;
    private final OllamaService ollamaService;
    private final FreemarkerService freemarkerService;
    private final UserFundRecordService userFundRecordService;

    /**
     * 将前端传递的生图参数封装为redis中的任务对象(包含了Comfyui生图的任务对象)
     * @param text2ImageReqDto 前端传递的生图参数
     * @return redis中的任务对象(包含了Comfyui生图的任务对象)
     */
    public ComfyuiTask getComfyuiTask(Text2ImageReqDto text2ImageReqDto) throws Exception {
        //1、将text2Image  ReqDto转换成ComfyuiModel
        //基本参数的封装
        ComfyuiModel model = BeanUtil.toBean (text2ImageReqDto, ComfyuiModel.class);
        model.setWidth (text2ImageReqDto.width ());
        model.setHeight (text2ImageReqDto.height ());
        model.setModelName (text2ImageReqDto.modelName ());
        model.setSamplerName (text2ImageReqDto.samplerName ());
        model.setScheduler (text2ImageReqDto.scheduler ());
        //提升词与负项提示词
        model.setPropmt (String.format("(8k, best quality, masterpiece),(high detailed skin),%s", ollamaService.translate(model.getPropmt())));
        if(StrUtil.isNotEmpty (model.getReverse ())){
            model.setReverse (ollamaService.translate (model.getReverse ()));
        }
        model.setReverse(String.format("%s,bad face,naked,bad finger,bad arm,bad leg,bad eye", model.getReverse()));

        //2、利用ComfyuiModel结合模板使用Freemarker渲染出comfyui json
        String prompt = freemarkerService.renderText2Image (model);

        //3、构建出Comfyui生图需要的对象ComfyuiRequestDto
        ComfyuiRequestDto dto = new ComfyuiRequestDto (COMFYUI_CLIENT_ID, JSON.parseObject (prompt));

        //4、构建出ComfyuiTask(Redis任务对象)
        ComfyuiTask task = new ComfyuiTask (text2ImageReqDto.getClientId (), dto);
        task.setSize (text2ImageReqDto.getSize ());
        return task;

    }

    @Override
    public Text2ImageResDto propmt(Text2ImageReqDto text2ImageReqDto) throws Exception {
        //积分冻结
        userFundRecordService.pointsFreeze (UserUtils.getUser ().getId (), text2ImageReqDto.getSize ());

        //1、获取Comfyui任务对象
        ComfyuiTask comfyuiTask = getComfyuiTask (text2ImageReqDto);

        //封装用户id
        comfyuiTask.setUserId (UserUtils.getUser ().getId ());

        //2、添加任务数据到redis中
        comfyuiTask = redisService.addQueueTask (comfyuiTask);

        //3、返回结果
        Text2ImageResDto resDto = new Text2ImageResDto ();
        resDto.setPid (comfyuiTask.getId ());
        resDto.setQueueIndex (comfyuiTask.getIndex ());
        return resDto;
    }
}
